#
#  This python script generates Fig.2 from the publication
#  It illustrates how to read in Python both complex output
#  of SPS and post-processed files corresponding to a
#  representation of the simulation fields in the physical space
#  as computed by for example by compute_fields_phys_space.m
#

# import required packages
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.axes_grid1 import make_axes_locatable

# Describe the simulation

Lx = 40.
Ly = 40.
Lz = 25.
nx = 399
ny = 399
nz = 249

# directories where that data is located

# physical representation (assumed to be single-precision)
# see compute_fields_phys_space.m
wdirR = "./data-physical/"
# double-precision comlex data
wdirS = "./data/" 

# time index to read
tidx = 1216

# maximum ky & kz to inlcude in spectral plot
kymax= 20
kzmax= 10

# define indices corresponding to k=0
ic = np.int(nx/2)
jc = np.int(ny/2)
kc = np.int(nz/2)
jmax = jc + np.int( kymax/2.0/np.pi*Ly)
kmax = kc + np.int( kzmax/2.0/np.pi*Lz)

# read physical space data for By (single precision float)
fname = "%s/by_%i.gda" % (wdirR,tidx)
b3D    = np.fromfile(fname,dtype=np.float32);
b3D    = np.reshape(b3D, [nx,ny,nz], order='F' );
# extract YZ-plane at X=Lx/2
b2D = np.squeeze(b3D[ic, :, :]).T

# read complex k-space data for By
fname = "%s/by_%i.bin" % (wdirS,tidx)
bk3D = np.fromfile(fname,dtype=np.float64);
# compute |By|^2 = (complex_part)^2 + (real_part)^2
S3Da  = bk3D[0::2]**2+bk3D[1::2]**2
# reshape |By|^2 into a 3D array
S3Da  = np.reshape(S3Da,[nx,ny,nz], order='F' );
# extract a plane at kx=0 with k in the range (0..kymax) and (0..kzmax)
S2Da  = np.squeeze(S3Da[ic, jc:jmax, kc:kmax]).T

# repeat the same steps for |Ey|^2
fname = "%s/ey_%i.bin" % (wdirS,tidx)
bk3D = np.fromfile(fname,dtype=np.float64);
S3Db  = bk3D[0::2]**2+bk3D[1::2]**2
S3Db  = np.reshape(S3Db,[nx,ny,nz], order='F' );
S2Db  = np.squeeze(S3Db[ic, jc:jmax, kc:kmax]).T

# choose color scheme
cmp = 'jet'

## ----------------------------------------------------
# make figure
## ----------------------------------------------------

f1 = plt.figure()
ax1 = f1.add_subplot( 3, 1, 1 )
ax3 = f1.add_subplot( 3, 1, 3 )
ax2 = f1.add_subplot( 3, 1, 2, sharex=ax3 )

# top panel: 2D plot of By in y-z plane at x=Lx/2
im1=ax1.imshow(b2D,origin='lower',extent=[0,Ly,0,Lz],vmin=-6E-5,
	       vmax=6E-5, cmap=cmp)
ax1.set_ylabel(r'$z/d_e$')
ax1.set_xlabel(r'$y/d_e$')
divider1 = make_axes_locatable(ax1)
cax1 = divider1.append_axes("right", size="5%", pad=0.05)
cbar1=f1.colorbar(im1, cax=cax1)
cbar1.formatter.set_powerlimits((-5, -6))
cbar1.update_ticks()

# middle panel: 2D plot of |By(kx=0,ky,kz)|^2 
im2=ax2.imshow(S2Da,origin='lower',extent=[0,kymax,0,kzmax],cmap=cmp)
ax2.set_ylabel(r'$k_{z}d_e$')
ax2.set_xlabel(r'$k_{y}d_e$')
divider2 = make_axes_locatable(ax2)
cax2 = divider2.append_axes("right", size="5%", pad=0.05)
f1.colorbar(im2, cax=cax2)

# bottom panel: 2D plot of |Ey(kx=0,ky,kz)|^2 
im3=ax3.imshow(S2Db,origin='lower',extent=[0,kymax,0,kzmax],cmap=cmp)
ax3.set_ylabel(r'$k_{z}d_e$')
ax3.set_xlabel(r'$k_{y}d_e$')
divider3 = make_axes_locatable(ax3)
cax3 = divider3.append_axes("right", size="5%", pad=0.05)
f1.colorbar(im3, cax=cax3)

# save figure
fname = "fig2.png"
plt.tight_layout(h_pad=0)	  
plt.savefig(fname,bbox_inches='tight')

    
